import requests
import hashlib
import zipfile
import os
import time

# --- 配置 ---
SYMBOL = "ETHUSDT"
BASE_URL_PREFIX = f"https://data.binance.vision/data/futures/um/monthly/klines/{SYMBOL}/"
TIME_INTERVALS = [
    "12h", "15m", "1d", "1h", "1m", "1mo", "1w", "2h",
    "30m", "3d", "3m", "4h", "5m", "6h", "8h"
]
START_DATE_STR = "2025-06"
END_DATE_STR = "2025-06"  # 包含此月份

# 请确保以下路径存在, 或者程序将尝试创建它们
# 使用原始字符串 (raw strings) 来处理Windows路径中的反斜杠
BASE_SAVE_PATH = r"F:\personal\binance_klines\\" + SYMBOL
ZIP_SAVE_DIR = os.path.join(BASE_SAVE_PATH, "zip")  # 所有zip文件保存到这里
EXTRACT_BASE_DIR = BASE_SAVE_PATH  # 解压文件的基础目录, 例如 F:\personal\binance_klines\12h\


# --- 辅助函数 ---

def download_file(url, save_path):
    """从URL下载文件到指定路径。"""
    try:
        print(f"尝试下载: {url}")
        # 确保目标文件夹存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)

        response = requests.get(url, stream=True, timeout=60)  # 设置超时时间
        response.raise_for_status()  # 如果HTTP请求返回了不成功的状态码, 则抛出HTTPError异常

        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"成功下载: {save_path}")
        return True
    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 404:
            print(f"文件未找到 (404): {url}")
        else:
            print(f"HTTP错误, 下载 {url} 失败: {e}")
    except requests.exceptions.RequestException as e:
        print(f"下载 {url} 错误: {e}")
    except Exception as e:
        print(f"下载 {url} 时发生未知错误: {e}")
    return False


def verify_checksum(zip_file_path, checksum_file_path):
    """对照CHECKSUM文件验证ZIP文件的SHA256校验和。"""
    if not os.path.exists(zip_file_path):
        print(f"用于校验和验证的ZIP文件不存在: {zip_file_path}")
        return False
    if not os.path.exists(checksum_file_path):
        print(f"CHECKSUM文件不存在: {checksum_file_path}")
        return False

    try:
        with open(checksum_file_path, 'r') as f:
            checksum_content = f.read().strip()

        # CHECKSUM文件通常格式为: "sha256_hash  filename.zip"
        # 我们只需要哈希值部分
        expected_hash = checksum_content.split()[0]

        sha256 = hashlib.sha256()
        with open(zip_file_path, 'rb') as f:
            while True:
                data = f.read(65536)  # 以64k块读取
                if not data:
                    break
                sha256.update(data)
        calculated_hash = sha256.hexdigest()

        if calculated_hash == expected_hash:
            print(f"文件 {os.path.basename(zip_file_path)} 校验和验证通过")
            return True
        else:
            print(f"文件 {os.path.basename(zip_file_path)} 校验和不匹配!")
            print(f"  期望值: {expected_hash}")
            print(f"  计算值: {calculated_hash}")
            return False
    except Exception as e:
        print(f"验证 {os.path.basename(zip_file_path)} 校验和时出错: {e}")
        return False


def unzip_file(zip_file_path, extract_to_dir):
    """将ZIP文件解压到指定目录。"""
    try:
        os.makedirs(extract_to_dir, exist_ok=True)
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to_dir)
        print(f"成功解压: {os.path.basename(zip_file_path)} 到 {extract_to_dir}")
        return True
    except zipfile.BadZipFile:
        print(f"错误: {os.path.basename(zip_file_path)} 是损坏的ZIP文件")
    except Exception as e:
        print(f"解压 {os.path.basename(zip_file_path)} 时出错: {e}")
    return False


# --- 主逻辑 ---
def main():
    # 创建基础保存目录 (如果不存在)
    os.makedirs(ZIP_SAVE_DIR, exist_ok=True)
    os.makedirs(EXTRACT_BASE_DIR, exist_ok=True)  # 例如 F:\personal\binance_klines

    start_year, start_month = map(int, START_DATE_STR.split('-'))
    end_year, end_month = map(int, END_DATE_STR.split('-'))

    current_year = start_year
    current_month = start_month

    while current_year < end_year or (current_year == end_year and current_month <= end_month):
        year_str = str(current_year)
        month_str = f"{current_month:02d}"  # 月份补零 (例如 1 -> 01)

        for interval in TIME_INTERVALS:
            zip_filename = f"{SYMBOL}-{interval}-{year_str}-{month_str}.zip"
            checksum_filename = f"{zip_filename}.CHECKSUM"

            # 构建下载链接
            # 示例: https://data.binance.vision/?prefix=data/futures/um/monthly/klines/ETHUSDT/12h/ETHUSDT-12h-2020-01.zip
            zip_url = f"{BASE_URL_PREFIX}{interval}/{zip_filename}"
            checksum_url = f"{BASE_URL_PREFIX}{interval}/{checksum_filename}"

            # 定义保存路径
            # ZIP文件保存路径: F:\personal\binance_klines\zip\ETHUSDT-12h-2020-01.zip
            final_zip_save_path = os.path.join(ZIP_SAVE_DIR, zip_filename)
            # CHECKSUM文件临时保存路径 (与ZIP文件同目录)
            checksum_save_path = os.path.join(ZIP_SAVE_DIR, checksum_filename)

            # 定义解压路径
            # 示例: F:\personal\binance_klines\12h\
            extract_dir_for_interval = os.path.join(EXTRACT_BASE_DIR, interval)

            print(f"\n--- 开始处理: {SYMBOL} {interval} {year_str}-{month_str} ---")

            # 检查是否已解压 (通常zip包内只有一个同名csv文件)
            # 假设解压后的csv文件名与zip包名类似 (去掉.zip后缀, 加上.csv后缀)
            expected_csv_filename = f"{SYMBOL}-{interval}-{year_str}-{month_str}.csv"
            expected_csv_path = os.path.join(extract_dir_for_interval, expected_csv_filename)

            if os.path.exists(expected_csv_path):
                print(f"数据已存在于 {expected_csv_path}, 跳过 {zip_filename}.")
                # 如果zip文件也存在, 可以选择性删除, 但按要求是保留zip的
                # if os.path.exists(final_zip_save_path):
                #     print(f"对应的ZIP文件 {final_zip_save_path} 也存在.")
                continue  # 跳到下一个文件处理

            # 如果ZIP文件存在但未解压, 尝试验证并解压
            if os.path.exists(final_zip_save_path) and not os.path.exists(expected_csv_path):
                print(f"ZIP文件 {final_zip_save_path} 已存在但未解压. 尝试验证和解压...")
                # 检查CHECKSUM文件是否存在, 不存在则下载
                if not os.path.exists(checksum_save_path):
                    print(f"对应的CHECKSUM文件 {checksum_save_path} 不存在, 尝试下载...")
                    if not download_file(checksum_url, checksum_save_path):
                        print(f"下载已存在ZIP的CHECKSUM文件失败. 跳过 {zip_filename}.")
                        time.sleep(1)  # 礼貌性等待
                        continue

                if verify_checksum(final_zip_save_path, checksum_save_path):
                    unzip_file(final_zip_save_path, extract_dir_for_interval)
                else:
                    print(f"已存在ZIP文件 {final_zip_save_path} 校验和失败. 建议删除并重新运行以下载。")
                    # 可选: os.remove(final_zip_save_path)

                # 验证后清理CHECKSUM文件
                if os.path.exists(checksum_save_path):
                    try:
                        os.remove(checksum_save_path)
                    except OSError as e:
                        print(f"无法删除CHECKSUM文件 {checksum_save_path}: {e}")
                continue  # 处理下一个文件

            # 1. 下载ZIP文件
            if not download_file(zip_url, final_zip_save_path):
                print(f"下载ZIP文件 {zip_filename} 失败, 跳过.")
                time.sleep(1)  # 发生错误时等待一下
                continue

            # 2. 下载CHECKSUM文件
            if not download_file(checksum_url, checksum_save_path):
                print(f"下载CHECKSUM文件 {checksum_filename} 失败. 无法验证 {zip_filename}.")
                # 可选: 如果校验和下载失败, 删除已下载的zip文件
                # if os.path.exists(final_zip_save_path):
                #     try:
                #         os.remove(final_zip_save_path)
                #         print(f"已删除不完整下载的ZIP: {final_zip_save_path}")
                #     except OSError as e:
                #         print(f"无法删除不完整下载的ZIP {final_zip_save_path}: {e}")
                time.sleep(1)
                continue

            # 3. 验证校验和
            if verify_checksum(final_zip_save_path, checksum_save_path):
                # 4. 如果校验通过, 则解压文件
                if not unzip_file(final_zip_save_path, extract_dir_for_interval):
                    print(f"文件 {zip_filename} 校验通过但解压失败.")
            else:
                print(f"文件 {zip_filename} 校验和验证失败. 不进行解压.")
                # 可选: 将校验失败的zip移动到特定文件夹或删除
                # failed_dir = os.path.join(ZIP_SAVE_DIR, "failed_checksum")
                # os.makedirs(failed_dir, exist_ok=True)
                # os.rename(final_zip_save_path, os.path.join(failed_dir, os.path.basename(final_zip_save_path)))

            # 5. 清理CHECKSUM文件 (已下载到ZIP_SAVE_DIR)
            # 要求是下载并验证, ZIP文件要求保留, CHECKSUM文件验证后可以删除
            if os.path.exists(checksum_save_path):
                try:
                    os.remove(checksum_save_path)
                    print(f"已删除CHECKSUM文件: {checksum_save_path}")
                except OSError as e:
                    print(f"无法删除CHECKSUM文件 {checksum_save_path}: {e}")

            time.sleep(0.5)  # 对服务器友好, 每次请求后稍作等待

        # 更新月份和年份
        current_month += 1
        if current_month > 12:
            current_month = 1
            current_year += 1
            # 如果年份已超过结束年份, 则停止 (即使月份可能还没到结束月份)
            # 这是为了确保如果END_DATE_STR是例如 "2025-04", 当current_year变成2026时循环能正确停止
            if current_year > end_year:
                break

    print("\n所有任务已完成.")


if __name__ == "__main__":
    main()
